import random
import math
import cPickle

class Agent(object):
    
    NAME = "Agent_007"
    Gamma = 0.9
    Alpha = 0.1
    num_of_actions = 4

    # states features:
    # [a] AmmoAmount    = 3 values [0:0, 1:1-2, 2:3+]
    # [b] TaskDistance  = 4 values [0:0, 1:1, 2:2, 3:3+]
    # Qtable[a][b][action] 
    AmmoAmount = 3
    TaskDistance = 4
    
    # Task = 3 values
    # Task = [0:CPShomebase, 1:CPSmidfield, 2:foebase]
    
    def initQtable(self):
        Qtable = [ [ [0 for i in range(self.num_of_actions)] \
                        for j in range(self.TaskDistance)] \
                        for k in range(self.AmmoAmount)]
        return Qtable
        
    def loadQtable(self):
        FILE = open('Qtable_agent2.txt','r')
        Qtable = cPickle.load(FILE)
        FILE.close()
        print "Qtable loaded"
        return Qtable
    
    def saveQtable(self):
        FILE = open('Qtable_agent2.txt','w')
        cPickle.dump(self.__class__.Qtable,FILE)
        FILE.close()
    
    def getState(self, o, t, a):
        # o = observationlist
        # t = task of agent
        # a = ammo locations

        #Determine AmmoAmount = 3 states
        if (o.ammo == 0):
            AmmoState = 0
        elif (1 <= o.ammo <= 2):
            AmmoState = 1
        else: # AmmoAmount >= 3
            AmmoState = 2
        
        #Determine TaskDistance = 4 states    
        s_pos = o.loc # location of agent
        t_pos = [0,0]    # location of task
        
        # get location of task
        if (t == 0):
            if (self.team == 0): #team = red
                t_pos = o.cps[0][0:2] #x,y position, CPS closest by red-team
            else:
                t_pos = o.cps[2][0:2]
        else: #(t == 1):
            t_pos = o.cps[1][0:2]
        
        # distance task-agent
        dx = t_pos[0]-s_pos[0]
        dy = t_pos[1]-s_pos[1]
        
        observationdistance = 100 #100px wide can the agent see
        
        # TaskID
        TaskState = int(math.sqrt(dx*dx + dy*dy) / observationdistance)
        if TaskState > 3:
            TaskState = 3
        
        return [AmmoState, TaskState]
    
    def getValue(self,state,action):
        value = self.__class__.Qtable[state[0]][state[1]][action]
        return value

    def setValue(self,state,action,value,reward):
        old_value = self.getValue(state,action)
        new_value = old_value + self.Alpha*(reward + self.Gamma*value - old_value)
        self.__class__.Qtable[state[0]][state[1]][action] = new_value

    def takeAction(self,a,obs):
        shoot = False
        if a == 0:
            # stay put
            goal = obs.loc
        elif a == 1:
            #go to task
            goal = obs.cps[self.__class__.task[self.id]][0:2]
        elif a == 2:
            # Move to closest Ammo-Location
            pos = obs.loc
            goal = None
            dist = 1000
            for loc in self.__class__.ammo_loc:
                dx = loc[0] - pos[0]
                dy = loc[1] - pos[1]
                d = math.sqrt(dx*dx + dy*dy)
                if d < dist:
                    goal = loc
                    dist = d
            if goal == None:
                goal = pos
        elif a == 3:
            # attack foe
            if (obs.foes):
                goal = obs.foes[0][0:2]
                if (obs.ammo > 0 and point_dist(obs.foes[0][0:2], obs.loc) < self.settings.max_range):
                       # and not line_intersects_grid(obs.loc, obs.foes[0][0:2], self.grid, self.settings.tilesize))
                # if you have ammo, you can reach a foe, and do not hit a friend
                    shoot = True
            else:
                #temporary goal to make the agent work
                goal = obs.cps[self.__class__.task[self.id]][0:2]
                #move to the closest agent of the opponent
        return (goal,shoot)
    
    def __init__(self, id, team, settings=None, field_rects=None, field_grid=None, nav_mesh=None):
        """ Each agent is initialized at the beginning of each game.
            The first agent (id==0) can use this to set up global variables.
            Note that the properties pertaining to the game field might not be
            given for each game.
        """
        self.id = id
        self.team = team
        self.mesh = nav_mesh
        self.grid = field_grid
        self.settings = settings
        self.goal = None
        self.old_state = []
        self.old_action = []
        self.old_cps = []
        
        # Recommended way to share variables between agents.
        if id == 0:
            self.__class__.Qtable = self.loadQtable()
            self.__class__.ammo_loc = []
            self.__class__.task = [0,0,1,1,1,1]
            
        #field size in pixels: (width, height)
        self.FIELDSIZE = (settings.tilesize*len(self.grid[0]), settings.tilesize*len(self.grid))

    
    def observe(self, observation):
        """ Each agent is passed an observation using this function,
            before being asked for an action. You can store either
            the observation object or its properties to use them
            to determine your action. Note that the observation object
            is modified in place.
        """
        self.observation = observation
        self.selected = observation.selected
        
        # Calculate the position of the Ammo-Locations
        ammopacks = filter(lambda x: x[2] == "Ammo", observation.objects)
        for pack in ammopacks:
            if pack[0:2] not in self.__class__.ammo_loc:
                self.__class__.ammo_loc.append(pack[0:2])
                # calculate position of mirrored version
                self.__class__.ammo_loc.append( (self.FIELDSIZE[0] - pack[0], pack[1]) )                
        
    def action(self):
        """ This function is called every step and should
            return a tuple in the form: (turn, speed, shoot)
        """
        obs = self.observation
        obs.cps.sort()
        state = self.getState(obs, self.__class__.task[self.id], obs.ammo)

        #choose best action
        best_val = -1000
        best_actions = []
        for a in range(self.num_of_actions):
            val = self.__class__.Qtable[state[0]][state[1]][a]#self.getValue(state,a)
            if val > best_val:
                best_val = val
                best_actions.append(a)
            if val == best_val:
                best_actions.append(a)
        best_action = best_actions[random.randint(0,len(best_actions)-1)]

        #Q update
        if self.old_state:
            new_cps_color = obs.cps[self.__class__.task[self.id]][2]
            old_cps_color = self.old_cps[self.__class__.task[self.id]][2]
            if (new_cps_color is not old_cps_color):
                if new_cps_color is self.team:
                    reward = 10
                else:
                    reward = -10
            else:
                reward = 0
            self.setValue(self.old_state,self.old_action,best_val,reward) 

        self.old_state = state
        self.old_action = best_action
        self.old_cps = obs.cps

        [self.goal,shoot] = self.takeAction(best_action,obs)

        # Compute path, angle and drive
        path = find_path(obs.loc, self.goal, self.mesh, self.grid, self.settings.tilesize)
        if path:
            dx = path[0][0]-obs.loc[0]
            dy = path[0][1]-obs.loc[1]
            turn = angle_fix(math.atan2(dy, dx)-obs.angle)
            if turn > self.settings.max_turn or turn < -self.settings.max_turn:
                shoot = False
            speed = (dx**2 + dy**2)**0.5
        else:
            turn = 0
            speed = 0
        
        return (turn,speed,shoot)
        
    def debug(self, surface):
        """ Allows the agents to draw on the game UI,
            Refer to the pygame reference to see how you can
            draw on a pygame.surface. The given surface is
            not cleared automatically. Additionally, this
            function will only be called when the renderer is
            active, and it will only be called for the active team.
        """
        import pygame
        # First agent clears the screen
        if self.id == 0:
            surface.fill((0,0,0,0))
        # Selected agents draw their info
        if self.selected:
            if self.goal is not None:
                pygame.draw.line(surface,(0,0,0),self.observation.loc, self.goal)
        
    def finalize(self, interrupted=False):
        """ This function is called after the game ends, 
            either due to time/score limits, or due to an
            interrupt (CTRL+C) by the user. Use it to
            store any learned variables and write logs/reports.
        """
        self.saveQtable()
        
